#pragma once

#include <windows.h>
#include <wtsapi32.h>
#include <userenv.h>
#include <string>

#pragma comment(lib, "wtsapi32.lib")
#pragma comment(lib, "userenv.lib")

class SystemServiceCapture {
public:
    SystemServiceCapture();
    ~SystemServiceCapture();

    bool Initialize();
    void Cleanup();
    bool CaptureActiveSession(HBITMAP& hBitmap, int& width, int& height);
    bool CaptureSpecificSession(DWORD sessionId, HBITMAP& hBitmap, int& width, int& height);
    std::string GetLastError() const { return m_lastError; }

private:
    bool EnableAllPrivileges();
    bool CaptureSessionDesktop(DWORD sessionId, HBITMAP& hBitmap, int& width, int& height);
    bool SwitchToSessionDesktop(DWORD sessionId);
    bool CaptureCurrentDesktopContent(HBITMAP& hBitmap, int& width, int& height);
    HDESK GetSessionDesktop(DWORD sessionId);
    bool IsSessionLocked(DWORD sessionId);

private:
    std::string m_lastError;
    bool m_initialized;
    HANDLE m_hToken;
    HDESK m_originalDesktop;
};